clc;
clear;
close all;

addpath('autogenV2\')

% 1、验证DH参数 passed
%{
alpha = [pi/2, pi/2, -pi/2, -1.4835];
a = [0, 0, 0, 0];

q = [pi/2, pi/2, pi/2, 0];
d = [0, 0, 0.3972, 0];

for i=1:4
    T(:,:,i)=[
            cos(q(i)), -sin(q(i)), 0, a(i);
            sin(q(i)).*cos(alpha(i)), cos(q(i))*cos(alpha(i)), -sin(alpha(i)), -d(i)*sin(alpha(i));
            sin(q(i))*sin(alpha(i)), cos(q(i))*sin(alpha(i)), cos(alpha(i)), d(i)*cos(alpha(i));
            0,0,0,1];
end

% 在axis_0坐标系下某点坐标
P = [0.01906; 0.1285; 0.2385; 1];
T01=T(:,:,1); P_a1 = inv(T01)*P; % 计算该点在axis_1坐标系下的值
T02=T01*T(:,:,2); P_a2 = inv(T02)*P; % 计算该点在axis_2坐标系下的值
T03=T02*T(:,:,3); P_a3 = inv(T03)*P; % 计算该点在axis_3坐标系下的值
T04=T03*T(:,:,4); P_a4 = inv(T04)*P; % 计算该点在axis_4坐标系下的值
%}

% 2、生成雅可比计算函数
%{
a = [0, 0, 0, 0];
d = [0, 0, 0.3972, 0];

% codegen
q=sym('q_%d',[4, 1],'real');
alpha = sym('alpha_%d',[4, 1],'real');

for i=1:4
    T(:,:,i)=[
            cos(q(i)), -sin(q(i)), 0, a(i);
            sin(q(i)).*cos(alpha(i)), cos(q(i))*cos(alpha(i)), -sin(alpha(i)), -d(i)*sin(alpha(i));
            sin(q(i))*sin(alpha(i)), cos(q(i))*sin(alpha(i)), cos(alpha(i)), d(i)*cos(alpha(i));
            0,0,0,1];
end

T01=T(:,:,1);
T02=T01*T(:,:,2);
T03=T02*T(:,:,3);
T04=T03*T(:,:,4);

for (i=1: 4)
    for (j=1: 4)
        % t1
        u41(i, j) = jacobian(T04(i, j), q(1));
        u31(i, j) = jacobian(T03(i, j), q(1));
        u21(i, j) = jacobian(T02(i, j), q(1));
        u11(i, j) = jacobian(T01(i, j), q(1));
        
        % t2
        u42(i, j) = jacobian(T04(i, j), q(2));
        u32(i, j) = jacobian(T03(i, j), q(2));
        u22(i, j) = jacobian(T02(i, j), q(2));
        
        % t3
        u43(i, j) = jacobian(T04(i, j), q(3));
        u33(i, j) = jacobian(T03(i, j), q(3));
        
        % t4
        u44(i, j) = jacobian(T04(i, j), q(4));
    end
end

optimize = true;
matlabFunction(u41, 'File', 'autogenV2/J_rl_T04_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u31, 'File', 'autogenV2/J_rl_T03_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u21, 'File', 'autogenV2/J_rl_T02_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u11, 'File', 'autogenV2/J_rl_T01_q1', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u42, 'File', 'autogenV2/J_rl_T04_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u32, 'File', 'autogenV2/J_rl_T03_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u22, 'File', 'autogenV2/J_rl_T02_q2', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u43, 'File', 'autogenV2/J_rl_T04_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u33, 'File', 'autogenV2/J_rl_T03_q3', 'Vars',{q, alpha}, 'Optimize', optimize);
matlabFunction(u44, 'File', 'autogenV2/J_rl_T04_q4', 'Vars',{q, alpha}, 'Optimize', optimize);
%}

% 3. 刚体重力矩验证

% CAD评估结果
m1 = 0.4269; 
r1_x=0.00002;r1_y=-0.0008;r1_z=-0.002;
m2 = 0.3108;
r2_x=0.0;r2_y=0.0894;r2_z=0.0001;
m3 = 0.2715;
r3_x=0.0;r3_y=-0.0002;r3_z=-0.0924;
m4 = 0.2283;
r4_x=-0.0045;r4_y=-0.2929;r4_z=0.0011;

% m1 = 0.527; 
% r1_x=0.00002;r1_y=-0.0008;r1_z=-0.002;
% m2 = 0.411;
% r2_x=0.0;r2_y=0.0894;r2_z=0.0001;
% m3 = 0.371;
% r3_x=0.0;r3_y=-0.0002;r3_z=-0.0924;
% m4 = 0.128;
% r4_x=-0.0045;r4_y=-0.2929;r4_z=0.0011;

% [0.526853101346153,0.410799924663641,0.371469906421152,0.128300109408913,
% 6.99999999563850e-05,-0.00134106439376176,-0.00200000000002220,
% -0.00499999999997764,0.0394386163405712,0.00413445582263687,
% -0.00499999603439588,-0.00383576235085579,-0.142299407267559,
% -0.00350000000007533,-0.242900000000022,0.000114558629156735,
% 4.55967632982577,3.88510278307024,4.34971172566065,4.37658612147982]

% 定义优化问题的上下界
lb = [
    m1-0.1, m2-0.1, m3-0.1, m4-0.1,...
    r1_x-0.005, r1_y-0.005, r1_z-0.005,...
    r2_x-0.005, r2_y-0.05, r2_z-0.0005,...
    r3_x-0.005, r3_y-0.005, r3_z-0.05,...
    r4_x-0.005, r4_y-0.05, r4_z-0.001, ...
    0, 0, 0, 0];

ub = [
    m1+0.1, m2+0.1, m3+0.1, m4+0.1,...
    r1_x+0.00005, r1_y+0.0, r1_z+0.0,...
    r2_x+0.005, r2_y+0.05, r2_z+0.05,...
    r3_x+0.005, r3_y+0.005, r3_z+0.05,...
    r4_x+0.001, r4_y+0.05, r4_z+0.005, ...
    5, 5, 5, 5];

m1_r1 = m1*[r1_x;r1_y;r1_z;1];
m2_r2 = m2*[r2_x;r2_y;r2_z;1];
m3_r3 = m3*[r3_x;r3_y;r3_z;1];
m4_r4 = m4*[r4_x;r4_y;r4_z;1];

data = load('dataset_v2/r_leg_01.txt');
data = data(200:505, :);

alpha = [pi/2, pi/2, -pi/2, -1.4835]';
% 提取其中的位置及力矩量
qs = data(:, 1:4);
tor = data(:, 9:12);
tor_est = data(:, 13:16);

g = [0 0 -9.8 0]';
tau_1s_real = [];tau_2s_real = [];tau_3s_real = [];tau_4s_real = [];
tau_1s_est_sw = [];tau_2s_est_sw = [];tau_3s_est_sw = [];tau_4s_est_sw = [];
Tau = [];
for i=1:1:length(qs)
    q = qs(i, :)'; 
    
    u11 = J_rl_T01_q1(q, alpha);u21 = J_rl_T02_q1(q, alpha);u31 = J_rl_T03_q1(q, alpha);u41 = J_rl_T04_q1(q, alpha);
    u22 = J_rl_T02_q2(q, alpha);u32 = J_rl_T03_q2(q, alpha);u42 = J_rl_T04_q2(q, alpha);
    u33 = J_rl_T03_q3(q, alpha);u43 = J_rl_T04_q3(q, alpha);
    u44 = J_rl_T04_q4(q, alpha);
    
    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4;
    tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4;
    tau_4 = -g'*u44*m4_r4;
    
    tau_4s_real = [tau_4s_real; tor(i, 4)];
    tau_3s_real = [tau_3s_real; tor(i, 3)];
    tau_2s_real = [tau_2s_real; tor(i, 2)];
    tau_1s_real = [tau_1s_real; tor(i, 1)];
    
    Tau = [Tau; 
    tor(i, 1)'; 
    tor(i, 2)'; 
    tor(i, 3)'; 
    tor(i, 4)'
    ];
    
    tau_4s_est_sw = [tau_4s_est_sw; tau_4*4.4];
    tau_3s_est_sw = [tau_3s_est_sw; tau_3*4.4];
    tau_2s_est_sw = [tau_2s_est_sw; tau_2*3.8];
    tau_1s_est_sw = [tau_1s_est_sw; tau_1*4.5];
end 

% subplot(2, 2, 1);
% plot(tau_1s_real, "r-");
% hold on;
% plot(tau_1s_est_sw, "b-");
% subplot(2, 2, 2);
% plot(tau_2s_real, "r.-");
% hold on;
% plot(tau_2s_est_sw, "b-");
% subplot(2, 2, 3);
% plot(tau_3s_real, "r.-");
% hold on;
% plot(tau_3s_est_sw, "b-");
% subplot(2, 2, 4);
% plot(tau_4s_real, "r.-");
% hold on;
% plot(tau_4s_est_sw, "b-");
% legend('real', 'est_{sw}');   

% 4.非线性最小二乘估计，估计所有的参数
x_init = [
    m1, m2, m3, m4, ... 
    r1_x, r1_y, r1_z, ...
    r2_x, r2_y, r2_z, ...
    r3_x, r3_y, r3_z, ...
    r4_x, r4_y, r4_z, ...
    3.5, 3.0, 2.0, 2.0]; 

% 只估计关节1的相关参数
% x_init = [m1, r1_x, r1_y, r1_z, 3.5];

qs_input = qs;
Taus = tau_estimator(x_init, qs);
x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau, lb, ub);
% x = lsqcurvefit(@tau_estimator, x_init, qs_input, Tau);

% CAD评估结果
% m1 = 0.4287; 
% r1_x=-0.01298;r1_y=-0.01016;r1_z=-0.000858;
% m2 = 0.3108;
% r2_x=0.0;r2_y=0.0894;r2_z=0.0001;
% m3 = 0.2715;
% r3_x=0.0;r3_y=-0.0002;r3_z=-0.0924;
% m4 = 0.2283;
% r4_x=-0.0045;r4_y=-0.2929;r4_z=0.0011;

% [0.794333962933533,0.621589663182557,0.385762322338198,0.165626046528936,
% 0.0319281011699424,-0.00512701907412078,-0.000858000000000000,
% -0.0328035904913637,0.106180935820085,-0.00495289448389376,
% -0.0100293028427989,-0.0173932188538317,-0.0756943377398151,
% 0.0106017175970497,-0.273490489353130,-0.0102697025270899,
% 3.20312709422456,2.57790923845779,2,3.01817593637323]

% [0.430660693167713,-0.0455916982009173,-0.0383585060500601,-0.000858000000000000,3.21388015592854]
% [0.430758880396943,-0.0518621117991189,-0.0307312389671013,-0.000858000000000000,3.14605186504073]
% [0.434772574590022,-0.103360078786451,-0.0235668748729218,-0.000858000000000000,2.80767360083987]

% m1_r1 = [x(1)*x(2); x(1)*x(3); x(1)*x(4); x(1)];  % 只估计关节1的参数

m1_r1 = [x(1)*x(5); x(1)*x(6); x(1)*x(7); x(1)];
m2_r2 = [x(2)*x(8); x(2)*x(9); x(2)*x(10); x(2)]; 
m3_r3 = [x(3)*x(11); x(3)*x(12); x(3)*x(13); x(3)]; 
m4_r4 = [x(4)*x(14); x(4)*x(15); x(4)*x(16); x(4)];

tau_1s_real = [];tau_2s_real = [];tau_3s_real = [];tau_4s_real = [];
tau_1s_est = [];tau_2s_est = [];tau_3s_est = [];tau_4s_est = [];
for i=1:1:length(qs_input)
    q = qs_input(i, :)'; 

    u11 = J_rl_T01_q1(q, alpha);u21 = J_rl_T02_q1(q, alpha);u31 = J_rl_T03_q1(q, alpha);u41 = J_rl_T04_q1(q, alpha);
    u22 = J_rl_T02_q2(q, alpha);u32 = J_rl_T03_q2(q, alpha);u42 = J_rl_T04_q2(q, alpha);
    u33 = J_rl_T03_q3(q, alpha);u43 = J_rl_T04_q3(q, alpha);
    u44 = J_rl_T04_q4(q, alpha);

    tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4;
    tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4;
    tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4;
    tau_4 = -g'*u44*m4_r4;

    tau_4s_real = [tau_4s_real; tor(i, 4)];
    tau_3s_real = [tau_3s_real; tor(i, 3)];
    tau_2s_real = [tau_2s_real; tor(i, 2)];
    tau_1s_real = [tau_1s_real; tor(i, 1)];
    
%     tau_1s_est = [tau_1s_est; tau_1*x(5)];  % 只估计关节1的参数
    tau_4s_est = [tau_4s_est; tau_4*x(20)];
    tau_3s_est = [tau_3s_est; tau_3*x(19)];
    tau_2s_est = [tau_2s_est; tau_2*x(18)];
    tau_1s_est = [tau_1s_est; tau_1*x(17)];
end

% plot
% hold on;
% plot(tau_1s_real, "r.-");
% hold on;
% plot(tau_1s_est, "g.-");
% legend('est_{sw}', 'real', 'est_{opt}');  

subplot(2, 2, 1);
plot(tau_1s_real, "r-");
hold on;
plot(tau_1s_est, "g-");
hold on;
plot(tau_1s_est_sw, "b-");
subplot(2, 2, 2);
plot(tau_2s_real, "r.-");
hold on;
plot(tau_2s_est, "g.-");
hold on;
plot(tau_2s_est_sw, "b-");
subplot(2, 2, 3);
plot(tau_3s_real, "r.-");
hold on;
plot(tau_3s_est, "g.-");
hold on;
plot(tau_3s_est_sw, "b-");
subplot(2, 2, 4);
plot(tau_4s_real, "r.-");
hold on;
plot(tau_4s_est, "g.-");
hold on;
plot(tau_4s_est_sw, "b-");
legend('real', 'est_{opt}', 'est_{sw}');   

%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function taus = tau_estimator(x, qs_input)
    alpha = [pi/2, pi/2, -pi/2, -1.4835]';
    g = [0 0 -9.8 0]';
    
    % 只估计关节1的参数
%     m1_r1 = [x(1)*x(2); x(1)*x(3); x(1)*x(4); x(1)];
    
    % 其他关节参数采用CAD估计的值
%     m2 = 0.3108;
%     r2_x=0.0;r2_y=0.0894;r2_z=0.0001;
%     m3 = 0.2715;
%     r3_x=0.0;r3_y=-0.0002;r3_z=-0.0924;
%     m4 = 0.2283;
%     r4_x=-0.0045;r4_y=-0.2929;r4_z=0.0011;
%     m2_r2 = m2*[r2_x;r2_y;r2_z;1];
%     m3_r3 = m3*[r3_x;r3_y;r3_z;1];
%     m4_r4 = m4*[r4_x;r4_y;r4_z;1];
    
% 优化所有参数
    m1_r1 = [x(1)*x(5); x(1)*x(6); x(1)*x(7); x(1)];
    m2_r2 = [x(2)*x(8); x(2)*x(9); x(2)*x(10); x(2)];
    m3_r3 = [x(3)*x(11); x(3)*x(12); x(3)*x(13); x(3)];
    m4_r4 = [x(4)*x(14); x(4)*x(15); x(4)*x(16); x(4)];
    

    taus = [];
    for i=1:1:length(qs_input)
        q = qs_input(i, :)'; 

        u11 = J_rl_T01_q1(q, alpha);u21 = J_rl_T02_q1(q, alpha);u31 = J_rl_T03_q1(q, alpha);u41 = J_rl_T04_q1(q, alpha);
        u22 = J_rl_T02_q2(q, alpha);u32 = J_rl_T03_q2(q, alpha);u42 = J_rl_T04_q2(q, alpha);
        u33 = J_rl_T03_q3(q, alpha);u43 = J_rl_T04_q3(q, alpha);
        u44 = J_rl_T04_q4(q, alpha);
        
        tau_1 = -g'*u11*m1_r1-g'*u21*m2_r2-g'*u31*m3_r3-g'*u41*m4_r4;
        tau_2 = -g'*u22*m2_r2-g'*u32*m3_r3-g'*u42*m4_r4;
        tau_3 = -g'*u33*m3_r3-g'*u43*m4_r4;
        tau_4 = -g'*u44*m4_r4;
        
        taus = [
            taus; 
%             tau_1*x(5)   % 只估计关节1的参数

            % 估计所有的参数
            tau_1*x(17); 
            tau_2*x(18); 
            tau_3*x(19); 
            tau_4*x(20)
            ];
    end
end
